# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: broker\worker_manager\jobs_storage\__init__.py
__author__ = "Attila Gerendi (Sanyi)"
import logging, datetime
from sqlalchemy import select as sql_select
from sqlalchemy import or_ as sql_or
from sqlalchemy import func as sql_func
from db import Connection
from db.tables.scan_session_jobs import ScanSessionJobRow
from helpers.constants.scans import *

logger = logging.getLogger("broker.worker_manager.pg_db_storage")


class JobsStorage:
    def __init__(self, shard):
        self.shard = shard

    def get(self, scan_session_id, scanning_app):
        q = (
            sql_select((ScanSessionJobRow,))
            .where(ScanSessionJobRow.scan_session_id == scan_session_id)
            .where(ScanSessionJobRow.scanning_app == scanning_app)
        )
        with Connection(self.shard) as (db):
            job = db.execute(q).fetchone()
            if not job:
                return
            return dict(job)

    def get_job_statistics_by_scanning_app(self, scanning_app, time_delta=0):
        q_jobs_lagging = (
            sql_select((sql_func.count(ScanSessionJobRow.scan_session_id),))
            .where(ScanSessionJobRow.scanning_app == scanning_app)
            .where(
                ScanSessionJobRow.status.in_(
                    (ScanStatusTypes.QUEUED, ScanStatusTypes.RESUMING)
                )
            )
            .where(
                sql_or(
                    ScanSessionJobRow.start_deadline.is_(None),
                    sql_func.current_timestamp()
                    > ScanSessionJobRow.start_deadline
                    + (datetime.timedelta(seconds=time_delta)),
                )
            )
        )
        q_jobs_about_to_finish = (
            sql_select((sql_func.count(ScanSessionJobRow.scan_session_id),))
            .where(ScanSessionJobRow.scanning_app == scanning_app)
            .where(
                ScanSessionJobRow.status.in_(
                    (ScanStatusTypes.ABORTING, ScanStatusTypes.PAUSING)
                )
            )
        )
        with Connection(self.shard) as (db):
            jobs_pending = db.execute(q_jobs_lagging).scalar()
            jobs_about_to_finish = db.execute(q_jobs_about_to_finish).scalar()
        return dict(lagging=jobs_pending, about_to_finish=jobs_about_to_finish)

    def get_active_jobs_count_by_worker_id(self, worker_id):
        q = (
            sql_select((sql_func.count(ScanSessionJobRow.scan_session_id),))
            .where(
                ScanSessionJobRow.status.in_(
                    (
                        ScanStatusTypes.STARTING,
                        ScanStatusTypes.ABORTING,
                        ScanStatusTypes.PAUSING,
                        ScanStatusTypes.PROCESSING,
                    )
                )
            )
            .where(ScanSessionJobRow.worker_id == worker_id)
        )
        with Connection(self.shard) as (db):
            r = db.execute(q).scalar()
        if r is None:
            return 0
        else:
            return r
